#include <vector>
#include <glm/glm.hpp>
#include <GL/glfw.h>
#include "Field.h"
#include "SimulationManager.h"

#define WINDOW_W 600
#define WINDOW_H 600

float ts = 0.03;

//#define TEST_MODE

//
// CELL INFO
#ifndef TEST_MODE
#define NUMCELLS 100

#else
#define NUMCELLS 100
std::vector<glm::vec2> testPartsPos;
std::vector<glm::vec2> testPartsVel;

#endif


// 
// PARTS INFO
#ifndef TEST_MODE
int NUMPARTSU = 300;
int NUMPARTSV = 300;
int NUMPARTS = NUMPARTSU*NUMPARTSV;
#else
int NUMPARTSU = 0, NUMPARTSV = 0, NUMPARTS = 0;
#endif

//
// Flags
bool doDrawGrid = false;
bool doDrawParticles = true;
bool doDrawVelocities = false;
bool doDrawDFs = false;
bool isPaused = false;
bool doStream = true;
bool doCollide = true;

#define VISC 0.8f // 0 - 2



float weight[9] = {
	4.0f/9.0f,

	1.0f/9.0f,
	1.0f/9.0f,
	1.0f/9.0f,
	1.0f/9.0f,

	1.0f/36.0f,
	1.0f/36.0f,
	1.0f/36.0f,
	1.0f/36.0f
};


glm::vec2 e[9] = {
	glm::vec2(0,0),

	glm::vec2(0,1),
	glm::vec2(1,0),
	glm::vec2(0,-1),
	glm::vec2(-1,0),

	glm::vec2(1,1),
	glm::vec2(1,-1),
	glm::vec2(-1,-1),
	glm::vec2(-1,1)
};

std::vector<glm::vec2> parts;
std::vector<glm::vec2> partsvel;

int currNdx = 0; // Current lb index to use.
float lbfield[2][NUMCELLS*NUMCELLS][9];


float lbfieldcntr[NUMCELLS*NUMCELLS] = {0.0f};
glm::vec2 lbfieldcentroid[NUMCELLS*NUMCELLS];
bool lbFieldObs[NUMCELLS*NUMCELLS] = {false};
float lbfielddens[2][NUMCELLS*NUMCELLS];
float amtRows = NUMCELLS*9;

float maxDens = std::numeric_limits<float>::min();
float minDens = std::numeric_limits<float>::max();
bool PERIODIC = false;



struct MouseState{ 
	MouseState() { isPressed = false; }
	glm::ivec2 last;
	bool isPressed;
} mouseState;

glm::ivec2 GetMouseCell(int mx, int my){
	glm::vec2 prop = glm::vec2( (float)mx / (float) WINDOW_W,
		(float)my / (float)WINDOW_H);

	return glm::ivec2(prop.x * NUMCELLS, prop.y * NUMCELLS);
}

#ifndef TEST_MODE
float supportRad = 4;
#else
float supportRad = 4;
#endif

int GetCellNdx(int cellU, int cellV){
	return cellV * NUMCELLS + cellU;
}

void GetT( const glm::vec2 & pos, glm::vec2 & t, int * cNdx){
// Get t value.
	int cellU = (int)pos.x,
		cellV = (int)pos.y;

	float d[4];
	float ptx = pos.x - (int)pos.x;
	float pty = pos.y - (int)pos.y;

	if (ptx >= 0.5f &&
		pty >= 0.5f)
	{
		t.x = ptx - 0.5f;
		t.y = pty - 0.5f;

		cNdx[0] = GetCellNdx(cellU, cellV);
		cNdx[1] = GetCellNdx(cellU+1, cellV);
		cNdx[2] = GetCellNdx(cellU, cellV+1);

		cNdx[3] = GetCellNdx(cellU+1, cellV+1);
	}
	else if (ptx < 0.5f &&
		pty >= 0.5f)
	{
		t.x = ptx + 0.5f;
		t.y = pty - 0.5f;

		cNdx[0] = GetCellNdx(cellU-1, cellV);
		cNdx[1] = GetCellNdx(cellU, cellV);
		cNdx[2] = GetCellNdx(cellU-1, cellV+1);
		cNdx[3] = GetCellNdx(cellU, cellV+1);
	}
	else if (ptx >= 0.5f &&
			 pty < 0.5f)
	{
		t.x = ptx - 0.5f;
		t.y = pty + 0.5f;

		cNdx[0] = GetCellNdx(cellU, cellV-1);
		cNdx[1] = GetCellNdx(cellU+1, cellV-1);
		cNdx[2] = GetCellNdx(cellU, cellV);
		cNdx[3] = GetCellNdx(cellU+1, cellV);
	}
	else if (pos.x - (int)pos.x < 0.5f &&
		pos.y - (int)pos.y < 0.5f)
	{
		t.x = (pos.x - (int)pos.x) + 0.5f;
		t.y = (pos.y - (int)pos.y) + 0.5f;

		cNdx[0] = GetCellNdx(cellU-1, cellV-1);
		cNdx[1] = GetCellNdx(cellU, cellV-1);
		cNdx[2] = GetCellNdx(cellU-1, cellV);
		cNdx[3] = GetCellNdx(cellU, cellV);
	}

}


void AddVelocity(glm::ivec2 & cell, glm::vec2 & inVel){
	float velMag = glm::length(inVel);
	glm::ivec2 tmp;
	for (int u = -supportRad; u < supportRad+1; ++u){
		for (int v = -supportRad; v < supportRad+1; ++v){
			tmp = glm::ivec2(cell.x+u, cell.y+v);
			int cndx = tmp.y*NUMCELLS + tmp.x;
			for (int i = 0; i < 9; ++i){
				float dot = glm::dot((e[i]), inVel);
				lbfield[currNdx][cndx][i] += dot*velMag*(velMag/(supportRad*2));
			}
		}
	}
}

void AddDensity(glm::ivec2 & cell){
	glm::ivec2 tmp;
	for (int u = -supportRad; u < supportRad; ++u){
		for (int v = -supportRad; v < supportRad; ++v){
			tmp = glm::ivec2(cell.x+u, cell.y+v);
			int cndx = tmp.y*NUMCELLS + tmp.x;
			lbfielddens[currNdx][cndx] = 1.0f;
		}
	}
}

void ProcessMouse () {
	glm::ivec2 mp;
	glfwGetMousePos(&mp.x, &mp.y);

	if (glfwGetMouseButton( GLFW_MOUSE_BUTTON_LEFT) == GLFW_PRESS){
		glm::ivec2 cell = GetMouseCell(mp.x, WINDOW_H - mp.y);

		glm::vec2 vel = glm::vec2(
			(float)cell.x-mouseState.last.x,
			(float)cell.y-mouseState.last.y);
		float lenvel = glm::length(vel);

		if (lenvel != 0){
			vel /= glm::length(vel);
			vel *= 0.5f;
		

			if (mouseState.isPressed){
				AddVelocity(cell, vel*0.01f);
				AddDensity(cell);
			}
		}

		if (mouseState.last != cell){
			mouseState.last = cell;
		}

		mouseState.isPressed = true;
	} 

	else {
		mouseState.isPressed = false;
	}
}



//A global simulation manager.
//This takes care of most boiler-plate code.
//If you only want to draw something, you only need to do three things:
//1) Create a display func with the function signerature
//	void funcname()
//2) In main, call 
//	sm.SetCallbackDisplay(funcname)
//  where funcname is the name of your draw function.
//3) In main, after the above function call, call
//	sm.Run();
//  This will start running the simulation.
// 
//  For a list of the callbacks the user can set, look at the "SetCallback..." functions 
//  in the SimulationManager header file.
//  Also, the GetCamera() method will get the camera and allow you to call methods which move
//	around the scene or rotate, etc.
SimulationManager sm(WINDOW_W, WINDOW_H);

//Display
//This function will be called as fast as possible, unlike the Update func.
//Only GL display code should be put in this section.  The screen
//is automatically cleared to white in the InitFrames() func, and the 
//glfwSwapBuffers() is called by the default EndFrames() function, so you only
//have to worry about actual drawing here.

void DrawDens() {
	glBegin(GL_QUADS);

	for (int i = 0; i < NUMCELLS; ++i){
		for (int j = 0; j < NUMCELLS; ++j){
			int cndx = j*NUMCELLS + i;

			if (lbfielddens[cndx] != 0){
				glColor3f(lbfielddens[currNdx][cndx], lbfielddens[currNdx][cndx]*0.25, lbfielddens[currNdx][cndx]*0.25);
			

				glVertex2i(i,j);
				glVertex2i(i+1,j);

				glVertex2i(i+1,j+1);
				glVertex2i(i,j+1);
			}
		}
	}
	glEnd();
}


void DrawGrid() {
	glBegin(GL_LINES);
	for (int i = 0; i < NUMCELLS; ++i){
		for (int j = 0; j < NUMCELLS; ++j){
			glVertex2f(i,j);
			glVertex2f(i+1,j);

			glVertex2f(i,j);
			glVertex2f(i,j+1);
		}
	}
	glEnd();
}

void DrawParts() {
	glPointSize(1.0f);

	glBegin(GL_POINTS);
	for (int i = 0; i < parts.size(); ++i){
		glVertex2fv(&parts[i][0]);
	}
	glEnd();
}



float CalcDensity( int cndx ){

	float sum = 0.0f;
	for (int i = 0; i < 9; ++i){
		sum += lbfield[currNdx][cndx][i];
	}

	return sum;
}

//void CalcWeightedVel( int cellU, int cellV, glm::vec2 pos, glm::vec2 & outVel ){
//	outVel = glm::vec2(0,0);
//	// Get vect to center
//	int cndx = GetCellNdx(cellU, cellV);
//	glm::vec2 cent = glm::vec2( (cellU +(cellU+1))*0.5f, (cellV + (cellV+1))*0.5f);
//
//	glm::vec2 vect = pos - cent;
//
//	for (int i = 0; i < 9; ++i){
//		float dot = 1;
//		if ( (vect.x != 0 || vect.y != 0) && (e[i].x != 0 || e[i].y != 0) ){
//			dot = std::fabs(glm::dot(glm::normalize(e[i]), glm::normalize(vect)));
//		}
//
//		outVel.x += e[i].x*lbfield[currNdx][cndx][i]*dot;
//		outVel.y += e[i].y*lbfield[currNdx][cndx][i]*dot;	
//	}
//
//}

void CalcWeightedVel( int cndx, glm::vec2 & outVel ){

	outVel = glm::vec2(0.0,0.0);

	for (int i = 0; i < 9; ++i){
		outVel.x += e[i].x*lbfield[currNdx][cndx][i];
		outVel.y += e[i].y*lbfield[currNdx][cndx][i];	
	}
}

void CalcVel( int cndx, glm::vec2 & outVel ){

	outVel = glm::vec2(0.0,0.0);

	for (int i = 0; i < 9; ++i){
		outVel.x += e[i].x*lbfield[currNdx][cndx][i];
		outVel.y += e[i].y*lbfield[currNdx][cndx][i];	
	}
}


float GetDiv( int cellU, int cellV ) {
	int cndx = GetCellNdx(cellU, cellV);
	glm::vec2 velxp, velxn, velyp, velyn;

	CalcVel(GetCellNdx(cellU+1, cellV), velxp);
	CalcVel(GetCellNdx(cellU-1, cellV), velxn);
	CalcVel(GetCellNdx(cellU, cellV+1), velyp);
	CalcVel(GetCellNdx(cellU, cellV-1), velyn);

	float dfdx = (velxp.x - velxn.x)*0.5f;
	float dfdy = (velyp.y - velyn.y)*0.5f;

	return dfdx + dfdy;
		
}

void DebugDiv(){
	for (int i = 0; i < NUMCELLS; ++i){
		for (int j = 0; j < NUMCELLS; ++j){
			float div = GetDiv(i, j);

			if (std::fabs(div) > 0.01){ printf("Div not 0: %f\n", div); }
		}
	}
}

void DrawLBField() {
	glBegin(GL_QUADS);

	for (int i = 0; i < NUMCELLS; ++i){
		for (int j = 0; j < NUMCELLS; ++j){
			bool draw = false;
			float dens = 0.0f;
			for (int df = 0; df < 9; ++df){
				if (lbfield[currNdx][j*NUMCELLS+i][df] != 0.0f) { draw = true; }

				dens += lbfield[currNdx][j*NUMCELLS+i][df]; 
			}

			//if (dens > 0) { dens /= maxDens; }

			if (abs(dens) < 0.06){ continue; }

			glColor3f( dens > 0.0f ? dens : 0.0f, 
				0, 
				dens < 0.0f ? dens : 0.0f);

			if (draw){
				glVertex2i(i,j);
				glVertex2i(i+1,j);
				glVertex2i(i+1,j+1);
				glVertex2i(i,j+1);
			}
		}
	}
	glEnd();
}


void DrawVector (float posx, float posy, const glm::vec2 & inVector, float mag ){

	glVertex2f(posx, posy);
	glVertex2f((float)posx + inVector.x*mag, (float)posy + inVector.y*mag);
}

#ifdef TEST_MODE
void DrawTestParticle(){
	glPointSize(2.0f);

	glBegin(GL_POINTS);

	for (int i = 0; i < testPartsPos.size(); ++i){
		glVertex2f(testPartsPos[i].x, testPartsPos[i].y);
	}

	glEnd();

	glBegin(GL_LINES);
	for (int i = 0; i < testPartsPos.size(); ++i){
		DrawVector(testPartsPos[i].x, testPartsPos[i].y, testPartsVel[i], 10.0f);
	}
	glEnd();

}
#endif

void DrawVelocities() {
	glBegin(GL_LINES);

	glm::vec2 vel;
	for (int i = 0; i < NUMCELLS; ++i){
		for (int j = 0; j < NUMCELLS; ++j){

			CalcVel(GetCellNdx(i,j), vel);

			//printf("vel: %f %f\n", vel.x, vel.y);

			glVertex2f((float)i+0.5f,(float)j+0.5f);
			glVertex2f((float)i+0.5f+vel.x*10, (float)j+0.5f+vel.y*20);
		}
	}
	glEnd();
}


void DrawDFs() {
	glBegin(GL_LINES);

	glm::vec2 vel;
	for (int i = 0; i < NUMCELLS; ++i){
		for (int j = 0; j < NUMCELLS; ++j){
			int cndx = GetCellNdx(i, j);
			for (int df = 0; df < 9; ++df){
				DrawVector( i+0.5f, j+0.5f, glm::vec2(e[df].x, e[df].y) * lbfield[currNdx][cndx][df], 10.0f);
			}

		
		}
	}
	glEnd();	
}

void Display() {

	glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
	glClear(GL_COLOR_BUFFER_BIT);

	glMatrixMode(GL_PROJECTION);
	glLoadIdentity();
	glOrtho(-1, NUMCELLS+1, -1, NUMCELLS+1, 0.01, 10 );

	glMatrixMode(GL_MODELVIEW);
	glLoadIdentity();

	glTranslatef(0,0,-1);

	glColor3f(1.0,0.0,0.0);

	DrawLBField();	

	//DrawDens();

	if (doDrawGrid){
		glColor3f(1.0,0.0,0.0);
		DrawGrid();
	}


	if (doDrawParticles){
		glColor3f(0.0, 1.0, 1.0);
		DrawParts();
	}

	if (doDrawVelocities){
		glColor3f(0,1.0,0);
		DrawVelocities();
	}

	if (doDrawDFs){
		glColor3f(1.0f, 0.8f, 0.0f);
		DrawDFs();
	}

#ifdef TEST_MODE
	glColor3f(0.0f, 0.0f, 1.0f);
	DrawTestParticle();
#endif
}

//Update
//This function will be called every timestep (ts).  The timestep can be
//set via the SetTimeStep method in SimulationManager. Put all non-drawing code
//here.
void KeyboardFunc( int key, int action ){

}


//
//ALL LB functiosINFORMATION HERE

void StreamCell( int cellU, int cellV) {
	// Calculate what the next cell is.  
	int offcellU, offcellV;
	int cndx = cellV*NUMCELLS + cellU;

	for (int i = 0; i < 9; ++i){
		glm::ivec2 offset = glm::ivec2((int)e[i].x, (int)e[i].y);

		// Minus to go  backwards
		offcellU = cellU - offset.x;
		offcellV = cellV - offset.y;

		// periodic
		if (PERIODIC){
			if (offcellU >= NUMCELLS){ offcellU = 0; }
			if (offcellU < 0){ offcellU = NUMCELLS-1; }
			if (offcellV >= NUMCELLS) { offcellV = 0; }
			if (offcellV < 0){ offcellV = NUMCELLS-1; }
		} 

		int offndx = offcellV*NUMCELLS + offcellU;

		if (PERIODIC){
			lbfield[(currNdx+1) % 2][cndx][i]  = lbfield[currNdx][offndx][i];
		} else {
			// Dealing with boundary cell in fi direction
			
			if (lbFieldObs[offndx] && i != 0){
				lbfield[(currNdx+1)%2][cndx][i] = lbfield[currNdx][offndx][ ((i+1)%4)+ (i < 5 ? 1 : 5) ] ;
			} else {
				lbfield[(currNdx+1) % 2][cndx][i]  = lbfield[currNdx][offndx][i];				
			}
		}

		
	}

}

float CalcFeq(int dfndx, int cndx, float dens, glm::vec2 & vel){
	float edotu = glm::dot(vel, e[dfndx]);

	// Rest density of 1
	return 
		weight[dfndx]*
		(dens + 
		3 * edotu - 
		(3.0f/2.0f)*(glm::dot(vel,vel)) + 
		(9.0f/2.0f)*(std::pow(edotu, 2.0f) ) );
}

void CollideCell( int cellU, int cellV ){
	int cndx = cellV*NUMCELLS + cellU;

	float dens = CalcDensity(cndx);

	// Set as max
	if (dens > maxDens) { maxDens = dens; }
	if (dens < minDens) { minDens = dens; }

	// XXX Should we divide by density?

	glm::vec2 vel;

	// XXX Testing

	CalcVel(cndx, vel);
	//vel = (vel.x != 0.0f && vel.y != 0.0f) ? glm::normalize(lbfieldcentroid[cndx])*4.0f : glm::vec2(0.0f, 0.0f);
	for (int i = 0; i < 9; ++i){
		float feq = CalcFeq(i, cndx, dens, vel);

		lbfield[(currNdx+1)%2][cndx][i]  = 
			(1.0f-VISC)*lbfield[(currNdx+1)%2][cndx][i] + VISC*feq;

	}
}

// Stream step.
void Stream() {
	for (int fcu = 0; fcu < NUMCELLS; ++fcu){
		for (int fcv = 0; fcv < NUMCELLS; ++fcv){
			if (!lbFieldObs[fcv*NUMCELLS + fcu] || PERIODIC){
				StreamCell( fcu, fcv );
			}
		}
	}
}

void Collide() {

	maxDens = std::numeric_limits<float>::min();
	minDens = std::numeric_limits<float>::max();

	for (int fcu = 0; fcu < NUMCELLS; ++fcu){
		for (int fcv = 0; fcv < NUMCELLS; ++fcv){			
				CollideCell( fcu, fcv );
		}
	}	

}

//End LB Functions


void ClearField() {
	for (int fn = 0; fn < 2; ++fn){
		for (int fc = 0; fc < NUMCELLS*NUMCELLS; ++fc){
			lbfielddens[fn][fc] = 0.0f;
			for (int df = 0; df < 9; ++df){
				lbfield[fn][fc][df] = 0.0f;
			}
		}
	}
}

#ifdef TEST_MODE
void UpdateTestPart() {
	for (int i = 0; i < testPartsPos.size(); ++i){
		int cellU = std::max(std::min(std::floor(testPartsPos[i].x), (float)NUMCELLS-1), 0.0f);
		int cellV = std::max(std::min(std::floor(testPartsPos[i].y), (float)NUMCELLS-1), 0.0f);
		int cndx = GetCellNdx(cellU, cellV);
		float dens = CalcDensity(cndx);
	
		glm::vec2 vel;
		CalcVel(cndx, vel);



		//if (partsvel[i].x != 0 || partsvel[i].y != 0){
		//	partsvel[i] = glm::normalize(partsvel[i]);
		//}
		//partsvel[i].y += -0.01f;

		testPartsVel[i] = vel;

		testPartsPos[i] += vel;

		//printf("Vel: %f, %f\n", testPartsVel[i].x, testPartsVel[i].y);
		//printf("Pos: %f, %f\n", testPartsPos[i].x, testPartsPos[i].y);
	}
}
#endif
void UpdateParts() {
	for (int i = 0; i < NUMPARTS; ++i){

		int cellU = std::max(std::min(std::floor(parts[i].x), (float)NUMCELLS-1), 0.0f);
		int cellV = std::max(std::min(std::floor(parts[i].y), (float)NUMCELLS-1), 0.0f);
		int cndx = GetCellNdx(cellU, cellV);
		float dens = CalcDensity(cndx);

		CalcVel(cndx, partsvel[i]);

		//if (partsvel[i].x != 0 || partsvel[i].y != 0){
		//	partsvel[i] = glm::normalize(partsvel[i]);
		//}
		//partsvel[i].y += -0.01f;

		parts[i] += partsvel[i];

		if ((int)parts[i].x < 0) { parts[i].x = 0; }
		if ((int)parts[i].x > NUMCELLS-1) { parts[i].x = NUMCELLS-1; }
		if ((int)parts[i].y < 0) { parts[i].y = 0; }
		if ((int)parts[i].y > NUMCELLS-1) { parts[i].y = NUMCELLS-1; }

	}
}

#ifdef TEST_MODE
#define NUMTESTPARTS 300

#endif


glm::vec2 lowerpc = glm::vec2(NUMCELLS*0.1f, NUMCELLS*0.1F);
glm::vec2 upperpc = glm::vec2(NUMCELLS*0.9f, NUMCELLS*0.9F);

#ifdef TEST_MODE
glm::vec2 stepsize = glm::vec2((upperpc-lowerpc).x/std::sqrt((float)NUMTESTPARTS), (upperpc-lowerpc).y/std::sqrt((float)NUMTESTPARTS));


#else
glm::vec2 stepsize = glm::vec2((upperpc-lowerpc).x/(float)NUMPARTSU, (upperpc-lowerpc).y/(float)NUMPARTSV);
#endif

void InitParts() {
#ifdef TEST_MODE
	testPartsPos.resize(NUMTESTPARTS);
	testPartsVel.resize(NUMTESTPARTS);

	int partnum = 0;
	for (float i = lowerpc.x; i < upperpc.x; i += stepsize.x){
		for (float j = lowerpc.y; j < upperpc.y; j += stepsize.y){
			if (partnum < NUMTESTPARTS){

				testPartsPos[partnum] = glm::vec2(i,j);

			}

			partnum++;
		}
	}

#else

	parts.resize(NUMPARTS);
	partsvel.resize(NUMPARTS);

	int partnum = 0;
	for (float i = lowerpc.x; i < upperpc.x; i += stepsize.x){
		for (float j = lowerpc.y; j < upperpc.y; j += stepsize.y){
			if (partnum < NUMPARTS){

				parts[partnum] = glm::vec2(i,j);
			}

			partnum++;
		}
	}
#endif
}

void InitTestField(){
	// set obstacle cells
	for (int i = 0; i < NUMCELLS; ++i){
		lbFieldObs[i] = true;
		lbFieldObs[i + (NUMCELLS-1)*NUMCELLS] = true;
	}
	for (int v = 0; v < NUMCELLS; ++v){
		lbFieldObs[v*NUMCELLS] = true;
		lbFieldObs[v*NUMCELLS + (NUMCELLS-1)] = true;
	}

	//for (int i = 30; i < 70; ++i){
	//	for (int j = 30; j < 70; ++j){
	//		int cndx = j*NUMCELLS + i;

	//		for (int i = 7; i < 9; ++i){
	//			lbfield[currNdx][cndx][i] = 0.25f;
	//		}
	//		lbfield[currNdx][cndx][4] = 0.25f;
	//		lbfield[currNdx][cndx][2] = 0.25f;
	//		
	//	}
	//}
}



void ClearLB() {
	for (int i = 0; i < NUMCELLS*NUMCELLS; ++i) {
		lbfieldcntr[i] = 0;
		lbfieldcentroid[i] = glm::vec2(0,0);

		/*for (int df = 0; df < 9; ++df){
			lbfield[currNdx][i][df] = 0.0f;

		}*/
	}
}



void Distribute(){

	// Put velocities into lattice
	for (int i = 0; i < NUMPARTS; ++i){
		glm::vec2 vel = partsvel[i];
		glm::vec2 pos = parts[i] - glm::vec2(std::floorf(parts[i].x)+0.5f, std::floorf(parts[i].y)+0.5f);
		int cellNdx = ((int)parts[i].y)*NUMCELLS + ((int)parts[i].x);
	
		for (int df = 0; df < 9; ++df){
			lbfield[currNdx][cellNdx][df] +=
				//std::max(0.0f, glm::dot(pos, e[df])) *//(std::sqrt(1.0f)-glm::length(pos)) *  //weighted based on particle.
				glm::dot(vel, e[df]); // amount of this velocity in that direction.

			lbfieldcntr[cellNdx] += 1.0f;
		}
	}

	//for (int i = 0; i < NUMCELLS*NUMCELLS; ++i){
	//	if (lbfieldcntr[i] != 0){
	//		for (int df = 0; df < 9; ++df){
	//			lbfield[currNdx][i][df] /= lbfieldcntr[i];
	//		}
	//	}
	//}
}


void SumVels() {

	for (int i = 0; i < NUMPARTS; ++i){
		glm::vec2 vel = partsvel[i];
		glm::vec2 pos = parts[i] - glm::vec2(std::floorf(parts[i].x)+0.5f, std::floorf(parts[i].y)+0.5f);
		int cellNdx = ((int)parts[i].y)*NUMCELLS + ((int)parts[i].x);
	

		lbfieldcentroid[cellNdx] += vel;
		lbfieldcntr[cellNdx] += 1.0f;
	}

	for (int i = 0; i < NUMCELLS*NUMCELLS; ++i){
		if (lbfieldcntr[i] != 0){
			for (int df = 0; df < 9; ++df){
				lbfieldcentroid[i] /= lbfieldcntr[i];
			}
		}
	}

}

float GetUserDens(int cellX, int cellY){
	if (cellX <= 0 || cellX >= NUMCELLS-1 || cellY <= 0 || cellY >= NUMCELLS-1){ return 0.0f; }
	return lbfielddens[currNdx][cellY*NUMCELLS+cellX];
}

float LinInterp( float * field, glm::vec2 pos){
	// Get t value.
	float tx, ty;
	int cellU = (int)pos.x,
		cellV = (int)pos.y;

	float d[4];
	float ptx = pos.x - (int)pos.x;
	float pty = pos.y - (int)pos.y;

	if (ptx >= 0.5f &&
		pty >= 0.5f)
	{
		tx = ptx - 0.5f;
		ty = pty - 0.5f;

		d[0] = GetUserDens(cellU, cellV);
		d[1] = GetUserDens(cellU+1, cellV);
		d[2] = GetUserDens(cellU, cellV+1);
		d[3] = GetUserDens(cellU+1, cellV+1);
	}
	else if (ptx < 0.5f &&
		pty >= 0.5f)
	{
		tx = ptx + 0.5f;
		ty = pty - 0.5f;

		d[0] = GetUserDens(cellU-1, cellV);
		d[1] = GetUserDens(cellU, cellV);
		d[2] = GetUserDens(cellU-1, cellV+1);
		d[3] = GetUserDens(cellU, cellV+1);
	}
	else if (ptx >= 0.5f &&
			 pty < 0.5f)
	{
		tx = ptx - 0.5f;
		ty = pty + 0.5f;

		d[0] = GetUserDens(cellU, cellV-1);
		d[1] = GetUserDens(cellU+1, cellV-1);
		d[2] = GetUserDens(cellU, cellV);
		d[3] = GetUserDens(cellU+1, cellV);
	}
	else if (pos.x - (int)pos.x < 0.5f &&
		pos.y - (int)pos.y < 0.5f)
	{
		tx = (pos.x - (int)pos.x) + 0.5f;
		ty = (pos.y - (int)pos.y) + 0.5f;

		d[0] = GetUserDens(cellU-1, cellV-1);
		d[1] = GetUserDens(cellU, cellV-1);
		d[2] = GetUserDens(cellU-1, cellV);
		d[3] = GetUserDens(cellU, cellV);
	}

//	printf("Tx ty: %f %f\n", tx, ty);

	return (1.0f - ty)*( (1.0f - tx)*d[0] + tx*d[1] ) +
				   ty* ( (1.0f - tx)*d[2] + tx*d[3] );
}

void Density(){

	for (int u = 0; u < NUMCELLS; ++u){
		for (int v = 0; v < NUMCELLS; ++v){
			int cndx = GetCellNdx(u, v);
			glm::vec2 vel;
			CalcVel(cndx, vel);
			// Advect density

			glm::vec2 pos = glm::vec2((float)u+0.5f, (float)v+0.5f) + (-1.0f * vel);

		
			float d = LinInterp(lbfielddens[(currNdx+1)%2], pos);
			lbfielddens[(currNdx+1)%2][cndx] = d;
		}
	}
}

#ifdef TEST_MODE
void DistributeTestParticleLin(){


	for (int i = 0; i < testPartsPos.size(); ++i){
		glm::vec2 vel = testPartsVel[i];
		glm::vec2 pos = testPartsPos[i];
		glm::vec2 posdiff = pos - 
			glm::vec2(std::floorf(pos.x)+0.5f, std::floorf(pos.y)+0.5f);

		glm::vec2 tuv;
		int cndxs[4];
		//ll, lr, ul, ur
		GetT(pos, tuv, cndxs);

		float linW = 0.0f;

		float cntr = 0;

		for (int i = 0; i < 4; ++i){

			if (i == 0){ linW = (1.0f-tuv.y)*(1.0f-tuv.x); }
			if (i == 1){ linW = (1.0f-tuv.y)*(tuv.x); }
			if (i == 2){ linW = (tuv.y)*(1.0f-tuv.x); }
			if (i == 3){ linW = tuv.y*tuv.x; }

			cntr += linW;

			for (int df = 0; df < 9; ++df){
				lbfield[currNdx][cndxs[i]][df] =
					//std::max(0.0f, glm::dot(pos, e[df])) *//(std::sqrt(1.0f)-glm::length(pos)) *  //weighted based on particle.
					glm::dot(vel, e[df])*linW; // amount of this velocity in that direction.
				//printf("df %d: %f\n", df, lbfield[currNdx][cndxs[i]][df]);
				lbfieldcntr[cndxs[i]] += 1.0f;
			}
		}

		
	}
}
#endif

#ifdef TEST_MODE
void DistributeTestParticle() {
	// Put velocities into lattice
	for (int i = 0; i < testPartsPos.size(); ++i){
		glm::vec2 vel = testPartsVel[i];
		glm::vec2 pos = testPartsPos[i];
		glm::vec2 posdiff = pos - 
			glm::vec2(std::floorf(pos.x)+0.5f, std::floorf(pos.y)+0.5f);
		


		int cellNdx = ((int)pos.y)*NUMCELLS + ((int)pos.x);
	
		for (int df = 0; df < 9; ++df){
			lbfield[currNdx][cellNdx][df] =
				//std::max(0.0f, glm::dot(pos, e[df])) *//(std::sqrt(1.0f)-glm::length(pos)) *  //weighted based on particle.
				glm::dot(vel, e[df]); // amount of this velocity in that direction.
			//printf("df %d: %f\n", df, lbfield[currNdx][cellNdx][df]);
			lbfieldcntr[cellNdx] += 1.0f;
		}
	}

}

#endif

void Update(float ts) {
	//XXX
	//All update code goes here
	//...
	// Set
	//ClearLB();

	if (!isPaused)
	{
#ifdef TEST_MODE
		//DistributeTestParticle();
		DistributeTestParticleLin();
#else
		//Distribute();
#endif
		//SumVels();

		ProcessMouse();
		// clear out previous LB cells

		if (doStream){
			Stream();
		}



		if (doCollide){
			Collide();
		}

		Density();
#ifdef TEST_MODE
	UpdateTestPart();
#else
		UpdateParts();
#endif
#ifdef TEST_MODE
		//isPaused = true;
#endif
	}
	//DebugDiv();

	currNdx = (currNdx+1)%2;
}

void Init(){
	glDisable(GL_DEPTH_TEST);
	
	ClearField();

	InitTestField();
	
	InitParts();

	sm.SetTimeStep(ts);
}

void AddVelocity() {
	for (int i = 30; i < 70; ++i){
		for (int j = 30; j < 70; ++j){
			int cndx = j*NUMCELLS + i;

			for (int i = 5; i < 6; ++i){
				lbfield[currNdx][cndx][i] = 0.25f;
			}
			//lbfield[currNdx][cndx][4] = 0.25f;
			//lbfield[currNdx][cndx][2] = 0.25f;
			
		}
	}
}

void DoDrawVelocity( ){
	doDrawVelocities = !doDrawVelocities;
}
void DoDrawParticles(){
	doDrawParticles = !doDrawParticles;
}

void DoDrawGrid() {
	doDrawGrid = !doDrawGrid;
}

void DoPause() {
	isPaused = !isPaused;
}

void DoCollide() {
	doCollide = !doCollide;
	printf("%d\n", doCollide);
}

void DoStream() {
	doStream = !doStream;
	printf("%d\n", doStream);
}

void DoDrawDFs(){
	doDrawDFs = !doDrawDFs;
}


void Test() {

	int cndx = 5*NUMCELLS + 5;
	//lbfield[currNdx][cndx][1] = 0.5f;
	lbfield[currNdx][cndx][5] = 0.5f;
	lbfield[currNdx][cndx][2] = 0.1f;


}

int main (int argc, char ** argv ){

#ifdef TEST_MODE
	isPaused = true;
#endif

	sm.SetCallbackDisplay(Display);
	sm.SetCallbackUpdate(Update);
	sm.SetCallbackLoadResources(Init);
	

	sm.AddKeyCallback('V', DoDrawVelocity);
	sm.AddKeyCallback('P', DoDrawParticles);
	sm.AddKeyCallback('G', DoDrawGrid);
	sm.AddKeyCallback('P', DoPause);
	sm.AddKeyCallback('C', DoCollide);
	sm.AddKeyCallback('S', DoStream);
	sm.AddKeyCallback('D', DoDrawDFs);

	sm.AddKeyCallback('T', Test);
	//XXX

	//Start simulation.
	sm.Run();

	return 0;
}